# 本程序使用ONNX模型，对指定路径的图片进行推理
# 在运行程序之前，要安装opencv-pyhton，numpy，onnxruntime
# 修改第11、12、13行，其它行无需修改

# 导入库
import cv2
import numpy as np
import onnxruntime as ort

# ---初始化设置---
model_path = 'model.onnx'   # 指定ONNX模型的位置
class_names = ['pen','udisk']   # 指定标签，与训练阶段的顺序一致
image_paths = ['test_pictures/pen.jpg','test_pictures/udisk.jpg']   # 指定所要检测的图片路径和对应标签
# ---设置完成，请运行程序---

# 创建results列表，用于存储多个推理结果
results = []

# 使用ONNX模型，对指定图片进行推理
for image_path in image_paths:
    # 图片预处理
    img = cv2.imread(image_path)
    img = cv2.resize(img,(224,224))
    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 根据模型表现，选择是否添加此语句
    img = np.expand_dims(img, axis=0)
    img = img.astype(np.float32)
    # 加载ONNX模型
    sess = ort.InferenceSession(model_path)
    input_name = sess.get_inputs()[0].name
    output_name = sess.get_outputs()[0].name
    # 推理并存储推理结果
    result = sess.run([output_name], {input_name: img})[0]
    results.append(result)

# 输出每一张图片的推理结果
for i in range(len(results)):
    result = results[i].tolist()
    print('Result of Inference:', result)
    max_index = result[0].index(max(result[0]))
    print(image_paths[i]," is ",class_names[max_index],'\n')
    
print('Done!')